#ifndef MYAMG_CSRMATRIX_CUH
#define MYAMG_CSRMATRIX_CUH
#include <cuda_runtime_api.h>
#include <driver_types.h>
#include <cstring>
#include <vector>
#include <string>
#include <sstream>
#include <iomanip>

namespace myamg {
    enum class MemoryType : unsigned char {
        Unknown,
        Host,
        Cuda,
        Unified
    };

    class MatrixHostAllocator {
    public:
        using SizeType = size_t;

        template<typename T>
        static T *allocate(SizeType n) {
            return new T[n];
        }

        template<typename T>
        static void deallocate(T *p) {
            delete[] p;
        }

        template<typename T>
        static void memcpy(T *dst, const T *src, SizeType n, cudaMemcpyKind kind) {
            std::memcpy(dst, src, sizeof(T) * n);
        }
    };

    class MatrixDeviceAllocator {
    public:
        using SizeType = size_t;

        template<typename T>
        static T *allocate(SizeType n) {
            T *ptr;
            cudaMalloc(&ptr, sizeof(T) * n);
            return ptr;
        }

        template<typename T>
        static void deallocate(T *p) {
            cudaFree(p);
        }

        template<typename T>
        static void memcpy(T *dst, const T *src, SizeType n, cudaMemcpyKind kind) {
            cudaMemcpy(dst, src, sizeof(T) * n, kind);
        }
    };

    template<typename Allocator>
    struct AllocatorType {
        static constexpr auto value = MemoryType::Unknown;
    };

    template<>
    struct AllocatorType<MatrixHostAllocator> {
        static constexpr auto value = MemoryType::Host;
    };

    template<>
    struct AllocatorType<MatrixDeviceAllocator> {
        static constexpr auto value = MemoryType::Cuda;
    };

    template<typename T>
    struct MemLoc {
        static constexpr auto type = MemoryType::Unknown;
    };

    template<typename Integer, typename Real, typename Allocator = MatrixHostAllocator>
    class CSRMatrix {
    public:
        using SizeType = size_t;

        CSRMatrix(SizeType nRows, SizeType nCols, SizeType nData)
            : m_nRows(nRows), m_nCols(nCols), m_nData(nData) {
            m_colIndex = Allocator::template allocate<Integer>(nData);
            m_rowPtr = Allocator::template allocate<Integer>(nRows + 1);
            m_values = Allocator::template allocate<Real>(nData);
        }

        CSRMatrix() : CSRMatrix(1, 1, 1) {}

        ~CSRMatrix() {
            Allocator::deallocate(m_colIndex);
            Allocator::deallocate(m_rowPtr);
            Allocator::deallocate(m_values);
        }

        CSRMatrix(const CSRMatrix &other)
            : m_nRows(other.m_nRows), m_nCols(other.m_nCols), m_nData(other.m_nData) {
            m_colIndex = Allocator::template allocate<Integer>(m_nData);
            m_rowPtr = Allocator::template allocate<Integer>(m_nRows + 1);
            m_values = Allocator::template allocate<Real>(m_nData);

            Allocator::template memcpy<Integer>(m_colIndex, other.m_colIndex, m_nData, cudaMemcpyDefault);
            Allocator::template memcpy<Integer>(m_rowPtr, other.m_rowPtr, m_nRows + 1, cudaMemcpyDefault);
            Allocator::template memcpy<Real>(m_values, other.m_values, m_nData, cudaMemcpyDefault);
        }

        CSRMatrix &operator=(const CSRMatrix &other) {
            if (this != &other) {
                if (m_nRows != other.m_nRows || m_nData != other.m_nData) {
                    Allocator::deallocate(m_colIndex);
                    Allocator::deallocate(m_rowPtr);
                    Allocator::deallocate(m_values);

                    m_nRows = other.m_nRows;
                    m_nCols = other.m_nCols;
                    m_nData = other.m_nData;

                    m_colIndex = Allocator::template allocate<Integer>(m_nData);
                    m_rowPtr = Allocator::template allocate<Integer>(m_nRows + 1);
                    m_values = Allocator::template allocate<Real>(m_nData);
                } else {
                    m_nCols = other.m_nCols;
                }

                Allocator::template memcpy<Integer>(m_colIndex, other.m_colIndex, m_nData, cudaMemcpyDefault);
                Allocator::template memcpy<Integer>(m_rowPtr, other.m_rowPtr, m_nRows + 1, cudaMemcpyDefault);
                Allocator::template memcpy<Real>(m_values, other.m_values, m_nData, cudaMemcpyDefault);
            }
            return *this;
        }

        CSRMatrix(CSRMatrix &&other) noexcept
            : m_colIndex(other.m_colIndex), m_rowPtr(other.m_rowPtr),
              m_values(other.m_values), m_nRows(other.m_nRows), m_nCols(other.m_nCols), m_nData(other.m_nData) {
            other.m_colIndex = nullptr;
            other.m_rowPtr = nullptr;
            other.m_values = nullptr;
            other.m_nRows = 0;
            other.m_nCols = 0;
            other.m_nData = 0;
        }

        CSRMatrix &operator=(CSRMatrix &&other) noexcept {
            if (this != &other) {
                Allocator::deallocate(m_colIndex);
                Allocator::deallocate(m_rowPtr);
                Allocator::deallocate(m_values);

                m_colIndex = other.m_colIndex;
                m_rowPtr = other.m_rowPtr;
                m_values = other.m_values;
                m_nRows = other.m_nRows;
                m_nCols = other.m_nCols;
                m_nData = other.m_nData;

                other.m_colIndex = nullptr;
                other.m_rowPtr = nullptr;
                other.m_values = nullptr;
                other.m_nRows = 0;
                other.m_nCols = 0;
                other.m_nData = 0;
            }
            return *this;
        }

        Integer *colIndex() { return m_colIndex; }
        const Integer *colIndex() const { return m_colIndex; }

        Integer *rowPtr() { return m_rowPtr; }
        const Integer *rowPtr() const { return m_rowPtr; }

        Real *values() { return m_values; }
        const Real *values() const { return m_values; }

        SizeType numRows() const { return m_nRows; }
        SizeType numCols() const { return m_nCols; }
        SizeType numNonzeros() const { return m_nData; }

        void setColIndex(const Integer *data, cudaMemcpyKind kind = cudaMemcpyDefault) {
            Allocator::template memcpy<Integer>(m_colIndex, data, m_nData, kind);
        }

        void setColIndex(const std::vector<Integer> &data) {
            Allocator::template memcpy<Integer>(m_colIndex, data.data(), m_nData, cudaMemcpyHostToDevice);
        }

        void setRowPtr(const Integer *data, cudaMemcpyKind kind = cudaMemcpyDefault) {
            Allocator::template memcpy<Integer>(m_rowPtr, data, m_nRows + 1, kind);
        }

        void setRowPtr(const std::vector<Integer> &data) {
            Allocator::template memcpy<Integer>(m_rowPtr, data.data(), m_nRows + 1, cudaMemcpyHostToDevice);
        }

        void setValues(const Real *data, cudaMemcpyKind kind = cudaMemcpyDefault) {
            Allocator::template memcpy<Real>(m_values, data, m_nData, kind);
        }

        void setValues(const std::vector<Real> &data) {
            Allocator::template memcpy<Real>(m_values, data.data(), m_nData, cudaMemcpyHostToDevice);
        }

        void setData(const Integer *colIndex, const Integer *rowPtr, const Real *values,
                     cudaMemcpyKind kind = cudaMemcpyDefault) {
            setColIndex(colIndex, kind);
            setRowPtr(rowPtr, kind);
            setValues(values, kind);
        }

        void setData(const std::vector<Integer> &colIndex, const std::vector<Integer> &rowPtr,
                     const std::vector<Real> &values) {
            setColIndex(colIndex);
            setRowPtr(rowPtr);
            setValues(values);
        }

        template<typename OtherAllocator>
        void copyFrom(const CSRMatrix<Integer, Real, OtherAllocator> &other) {
            Allocator::template memcpy<Integer>(m_colIndex, other.colIndex(), m_nData, cudaMemcpyDefault);
            Allocator::template memcpy<Integer>(m_rowPtr, other.rowPtr(), m_nRows + 1, cudaMemcpyDefault);
            Allocator::template memcpy<Real>(m_values, other.values(), m_nData, cudaMemcpyDefault);
        }

        void loadFromDense(const Real *denseData, SizeType nRows, SizeType nCols,
                           Real tolerance = Real(0), cudaMemcpyKind kind = cudaMemcpyDefault) {
            if (nCols != m_nCols) {
                return;
            }

            Real *hostData = nullptr;
            bool needFree = false;

            if (kind == cudaMemcpyDeviceToHost || kind == cudaMemcpyDeviceToDevice) {
                hostData = new Real[nRows * nCols];
                cudaMemcpy(hostData, denseData, sizeof(Real) * nRows * nCols, cudaMemcpyDeviceToHost);
                needFree = true;
            } else {
                hostData = const_cast<Real *>(denseData);
            }

            SizeType nnz = 0;
            for (SizeType i = 0; i < nRows * nCols; ++i) {
                if (hostData[i] > tolerance || hostData[i] < -tolerance) {
                    ++nnz;
                }
            }

            if (nnz != m_nData || nRows != m_nRows) {
                Allocator::deallocate(m_colIndex);
                Allocator::deallocate(m_rowPtr);
                Allocator::deallocate(m_values);

                m_nRows = nRows;
                m_nData = nnz;

                m_colIndex = Allocator::template allocate<Integer>(m_nData);
                m_rowPtr = Allocator::template allocate<Integer>(m_nRows + 1);
                m_values = Allocator::template allocate<Real>(m_nData);
            }

            Integer *hostRowPtr = new Integer[m_nRows + 1];
            Integer *hostColIndex = new Integer[m_nData];
            Real *hostValues = new Real[m_nData];

            SizeType idx = 0;
            hostRowPtr[0] = 0;
            for (SizeType i = 0; i < nRows; ++i) {
                for (SizeType j = 0; j < nCols; ++j) {
                    Real val = hostData[i * nCols + j];
                    if (val > tolerance || val < -tolerance) {
                        hostColIndex[idx] = j;
                        hostValues[idx] = val;
                        ++idx;
                    }
                }
                hostRowPtr[i + 1] = idx;
            }

            Allocator::template memcpy<Integer>(m_rowPtr, hostRowPtr, m_nRows + 1, cudaMemcpyHostToDevice);
            Allocator::template memcpy<Integer>(m_colIndex, hostColIndex, m_nData, cudaMemcpyHostToDevice);
            Allocator::template memcpy<Real>(m_values, hostValues, m_nData, cudaMemcpyHostToDevice);

            delete[] hostRowPtr;
            delete[] hostColIndex;
            delete[] hostValues;
            if (needFree) {
                delete[] hostData;
            }
        }

        void loadFromDense(const std::vector<Real> &denseData, SizeType nRows, SizeType nCols,
                           Real tolerance = Real(0)) {
            loadFromDense(denseData.data(), nRows, nCols, tolerance, cudaMemcpyHostToHost);
        }

        std::string toDenseString(int precision = 6) const {
            Integer *hostRowPtr = new Integer[m_nRows + 1];
            Integer *hostColIndex = new Integer[m_nData];
            Real *hostValues = new Real[m_nData];

            Allocator::template memcpy<Integer>(hostRowPtr, m_rowPtr, m_nRows + 1, cudaMemcpyDeviceToHost);
            Allocator::template memcpy<Integer>(hostColIndex, m_colIndex, m_nData, cudaMemcpyDeviceToHost);
            Allocator::template memcpy<Real>(hostValues, m_values, m_nData, cudaMemcpyDeviceToHost);

            std::vector<Real> dense(m_nRows * m_nCols, Real(0));
            for (SizeType i = 0; i < m_nRows; ++i) {
                for (Integer idx = hostRowPtr[i]; idx < hostRowPtr[i + 1]; ++idx) {
                    dense[i * m_nCols + hostColIndex[idx]] = hostValues[idx];
                }
            }

            std::ostringstream oss;
            oss << std::fixed << std::setprecision(precision);
            for (SizeType i = 0; i < m_nRows; ++i) {
                for (SizeType j = 0; j < m_nCols; ++j) {
                    oss << std::setw(precision + 4) << dense[i * m_nCols + j];
                    if (j < m_nCols - 1) {
                        oss << " ";
                    }
                }
                if (i < m_nRows - 1) {
                    oss << "\n";
                }
            }

            delete[] hostRowPtr;
            delete[] hostColIndex;
            delete[] hostValues;

            return oss.str();
        }

    private:
        Integer *m_colIndex;
        Integer *m_rowPtr;
        Real *m_values;
        SizeType m_nRows;
        SizeType m_nCols;
        SizeType m_nData;
    };

    template<typename Integer, typename Real, typename Allocator>
    struct MemLoc<CSRMatrix<Integer, Real, Allocator>> {
        static constexpr auto type = AllocatorType<Allocator>::value;
    };

    template<typename Real, typename Allocator = MatrixHostAllocator>

    class Vector {
    public:
        using SizeType = size_t;

        Vector(SizeType nRows)
            : m_nRows(nRows) {
            m_values = Allocator::template allocate<Real>(nRows);
        }

        Vector() : Vector(1) {}

        ~Vector() {
            Allocator::deallocate(m_values);
        }

        Vector(const Vector &other)
            : m_nRows(other.m_nRows) {
            m_values = Allocator::template allocate<Real>(m_nRows);
            Allocator::template memcpy<Real>(m_values, other.m_values, m_nRows, cudaMemcpyDefault);
        }

        Vector &operator=(const Vector &other) {
            if (this != &other) {
                if (m_nRows != other.m_nRows) {
                    Allocator::deallocate(m_values);
                    m_nRows = other.m_nRows;
                    m_values = Allocator::template allocate<Real>(m_nRows);
                }
                Allocator::template memcpy<Real>(m_values, other.m_values, m_nRows, cudaMemcpyDefault);
            }
            return *this;
        }

        Vector(Vector &&other) noexcept
            : m_values(other.m_values), m_nRows(other.m_nRows) {
            other.m_values = nullptr;
            other.m_nRows = 0;
        }

        Vector &operator=(Vector &&other) noexcept {
            if (this != &other) {
                Allocator::deallocate(m_values);
                m_values = other.m_values;
                m_nRows = other.m_nRows;
                other.m_values = nullptr;
                other.m_nRows = 0;
            }
            return *this;
        }

        Real *values() { return m_values; }
        const Real *values() const { return m_values; }

        SizeType size() const { return m_nRows; }

        void setValues(const Real *data, cudaMemcpyKind kind = cudaMemcpyDefault) {
            Allocator::template memcpy<Real>(m_values, data, m_nRows, kind);
        }

        void setValues(const std::vector<Real> &data) {
            Allocator::template memcpy<Real>(m_values, data.data(), m_nRows, cudaMemcpyHostToDevice);
        }

        template<typename OtherAllocator>
        void copyFrom(const Vector<Real, OtherAllocator> &other) {
            Allocator::template memcpy<Real>(m_values, other.values(), m_nRows, cudaMemcpyDefault);
        }

        std::string toString(int precision = 6) const {
            Real *hostValues = new Real[m_nRows];
            Allocator::template memcpy<Real>(hostValues, m_values, m_nRows, cudaMemcpyDeviceToHost);

            std::ostringstream oss;
            oss << std::fixed << std::setprecision(precision);
            for (SizeType i = 0; i < m_nRows; ++i) {
                oss << std::setw(precision + 4) << hostValues[i];
                if (i < m_nRows - 1) {
                    oss << "\n";
                }
            }

            delete[] hostValues;
            return oss.str();
        }

    private:
        Real *m_values;
        SizeType m_nRows;
    };

    template<typename Real, typename Allocator>
    struct MemLoc<Vector<Real, Allocator>> {
        static constexpr auto type = AllocatorType<Allocator>::value;
    };

    template<typename T>
    MemoryType getMemType(const T&) {
        return MemLoc<T>::type;
    }
}

#endif // MYAMG_CSRMATRIX_CUH
